import logging
import os

import streamlit as st
from llama_index.core import Settings, VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.chat_engine import CondenseQuestionChatEngine
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope

# 配置日志记录
logging.basicConfig(level=logging.DEBUG)

'''
streamlit run class03/p3_3_5.py
'''
@st.cache_resource(show_spinner=False)
def load_data():
    """加载文档数据并创建索引，使用缓存提高性能"""
    # 初始化模型对象
    llm_dash = DashScope(
        model="qwen-max",  # 使用通义千问Max模型
        api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从环境变量获取API密钥
        temperature=0.3,  # 控制生成结果的随机性
        system_prompt="""作为手机选购助手，我将帮助用户选择合适手机。
              仅根据提供的文档内容进行回答。
              如果问题与文档无关，请回答：抱歉，这个话题我不擅长。"""
    )

    # 使用本地嵌入模型（节省成本）
    embed_model_zh = HuggingFaceEmbedding(
        model_name="BAAI/bge-small-zh-v1.5",  # 中文优化的嵌入模型
        device="cpu"  # 使用CPU运行
    )

    # 显示加载进度提示
    with st.spinner(text="正在加载文档数据，请稍候..."):
        data_dir = 'D:/pythonRoom/llm_llama/data2'  # 文档数据目录
        reader = SimpleDirectoryReader(input_dir=data_dir, recursive=True)
        docs = reader.load_data()  # 加载文档数据

        # 配置全局设置
        Settings.embed_model = embed_model_zh
        Settings.llm = llm_dash

        # 创建向量索引
        index = VectorStoreIndex.from_documents(docs)
        return index


# 自定义CSS样式实现左右布局
def apply_custom_styles():
    """应用自定义CSS样式实现左右布局的聊天界面"""
    st.markdown(
        """
        <style>
        /* 用户消息样式 - 右侧 */
        .user-message {
            display: flex;
            justify-content: flex-end;
            margin-bottom: 15px;
        }

        .user-content {
            background-color: #dcf8c6;
            border-radius: 15px 15px 0 15px;
            padding: 10px 15px;
            max-width: 80%;
            box-shadow: 0 1px 2px rgba(0,0,0,0.1);
        }

        /* 助手消息样式 - 左侧 */
        .assistant-message {
            display: flex;
            justify-content: flex-start;
            margin-bottom: 15px;
        }

        .assistant-content {
            background-color: #ffffff;
            border-radius: 15px 15px 15px 0;
            padding: 10px 15px;
            max-width: 80%;
            box-shadow: 0 1px 2px rgba(0,0,0,0.1);
        }

        /* 头像样式 */
        .avatar {
            width: 40px;
            height: 40px;
            border-radius: 50%;
            display: flex;
            align-items: center;
            justify-content: center;
            margin: 0 10px;
            font-size: 20px;
        }

        .user-avatar {
            background-color: #128C7E;
            color: white;
        }

        .assistant-avatar {
            background-color: #25D366;
            color: white;
        }

        /* 消息容器 */
        .message-container {
            display: flex;
            align-items: flex-start;
            width: 100%;
        }

        /* 输入框样式 */
        .stChatInput {
            position: fixed;
            bottom: 20px;
            width: calc(100% - 40px);
            background: white;
            padding: 10px;
            border-radius: 10px;
            box-shadow: 0 -2px 10px rgba(0,0,0,0.1);
        }
        </style>
        """,
        unsafe_allow_html=True
    )


# 主程序入口
if __name__ == "__main__":
    # 设置页面标题和图标
    st.set_page_config(
        page_title="手机选购助手",
        page_icon="📱",
        layout="centered",
        initial_sidebar_state="collapsed"
    )

    # 应用自定义样式
    apply_custom_styles()

    # 初始化聊天引擎
    if "chat_engine" not in st.session_state:
        with st.spinner("正在初始化AI助手..."):
            # 加载数据并创建索引
            vector_index = load_data()
            # 创建查询引擎
            query_engine = vector_index.as_query_engine()
            # 创建聊天引擎（带上下文压缩功能）
            chat_engine = CondenseQuestionChatEngine.from_defaults(
                query_engine=query_engine,
                verbose=True  # 输出详细日志
            )
            st.session_state.chat_engine = chat_engine

    # 初始化聊天消息历史
    if "messages" not in st.session_state:
        st.session_state.messages = [
            {"role": "assistant", "content": "您好！我是手机选购助手，请问有什么可以帮您？"}
        ]

    # 创建聊天容器
    chat_container = st.container()

    # 在容器内显示所有历史聊天消息
    with chat_container:
        for message in st.session_state.messages:
            if message["role"] == "assistant":
                # 左侧显示助手消息
                st.markdown(
                    f"""
                    <div class="assistant-message">
                        <div class="message-container">
                            <div class="avatar assistant-avatar">🤖</div>
                            <div class="assistant-content">
                                {message["content"]}
                            </div>
                        </div>
                    </div>
                    """,
                    unsafe_allow_html=True
                )
            else:
                # 右侧显示用户消息
                st.markdown(
                    f"""
                    <div class="user-message">
                        <div class="message-container">
                            <div class="user-content">
                                {message["content"]}
                            </div>
                            <div class="avatar user-avatar">👤</div>
                        </div>
                    </div>
                    """,
                    unsafe_allow_html=True
                )

    # 获取用户输入
    if prompt := st.chat_input("请输入您的手机相关问题..."):
        # 添加用户消息到历史记录
        st.session_state.messages.append({"role": "user", "content": prompt})

        # 在容器内显示用户消息（右侧）
        with chat_container:
            st.markdown(
                f"""
                <div class="user-message">
                    <div class="message-container">
                        <div class="user-content">
                            {prompt}
                        </div>
                        <div class="avatar user-avatar">👤</div>
                    </div>
                </div>
                """,
                unsafe_allow_html=True
            )

        # 生成助手回复
        with chat_container:
            # 添加占位符用于显示思考状态
            thinking_placeholder = st.empty()
            thinking_placeholder.markdown(
                """
                <div class="assistant-message">
                    <div class="message-container">
                        <div class="avatar assistant-avatar">🤖</div>
                        <div class="assistant-content">
                            <i>正在思考...</i>
                        </div>
                    </div>
                </div>
                """,
                unsafe_allow_html=True
            )

            try:
                # 获取聊天引擎的回复
                response = st.session_state.chat_engine.chat(prompt)
                response_content = response.response
            except Exception as e:
                response_content = f"抱歉，处理您的请求时出错: {str(e)}"

            # 更新占位符为实际回复
            thinking_placeholder.markdown(
                f"""
                <div class="assistant-message">
                    <div class="message-container">
                        <div class="avatar assistant-avatar">🤖</div>
                        <div class="assistant-content">
                            {response_content}
                        </div>
                    </div>
                </div>
                """,
                unsafe_allow_html=True
            )

        # 添加助手回复到历史记录
        st.session_state.messages.append(
            {"role": "assistant", "content": response_content}
        )
